"""
Eyepacs 数据集预处理脚本
----------------------------------
本脚本用于对 EyePACS 及相关数据集进行预处理，完成以下工作：
1. 从 CSV 文件中读取 DR 分级信息，转化为文本诊断描述；
2. 对图像进行黑边裁剪、resize（默认 512x512），并保存为 PNG 格式；
3. 统一命名为 {prefix}_{原始文件名}.png；
4. 生成 annotations.json（标注信息） 与 split.json（数据集划分信息）；
5. 可选统计功能：统计分类任务的类别数量及数据规模。

使用示例：
-----------
# 命令行调用
python preprocess_eyepacs.py \
    --data_path ../../Dataset/public_dataset/EyePACS \
    --tar_path ../../Dataset/public_processed/EyePACS \
    --prefix eyepacs \
    --resize 512 512

输出文件：
-----------
- images/                # 处理后的图像
- annotations.json       # 图像标注信息
- split.json             # 数据集划分信息
- eyepacs.json (可选) # 分类任务统计信息（类别分布、总数等）
"""

import sys, os
import json
import csv
import numpy as np
from PIL import Image
from tqdm import tqdm
from data_preprocess.utils.crop import crop_resize_save
from data_preprocess.utils.stat import save_classification_stats
def parse_csv(csv_path):
    """
    读取 CSV 文件，并返回一个字典，格式如下：
    
    {
       image_id: "<diagnosis text>"
    }
    
    其中，文本描述的生成规则为：
      - 对于 adjudicated_dr_grade（或 level，ICDR 5 级）:
           0 -> normal
           1 -> Mild Diabetic Retinopathy
           2 -> Moderate Diabetic Retinopathy
           3 -> Severe Diabetic Retinopathy
           4 -> Proliferative Diabetic Retinopathy
           
      - 如果 CSV 文件中既没有 "level" 也没有 "DR_grade" 列，则直接抛出异常。
      
    示例 CSV 文件内容：
        ,image,quality,DR_grade
        0,1_right.jpeg,1,0
    """
    dr_mapping = {
        "0": "normal",
        "1": "Mild Diabetic Retinopathy",
        "2": "Moderate Diabetic Retinopathy",
        "3": "Severe Diabetic Retinopathy",
        "4": "Proliferative Diabetic Retinopathy"
    }
    quality_mapping = {
        "0":"Good"
        ,"1":"Usable"
        ,"2":"Reject"
    }
    result = {}
    with open(csv_path, newline='', encoding='utf-8') as f:
        reader = csv.DictReader(f)
        fieldnames = reader.fieldnames or []
        # 检查使用哪一列来表示 DR 等级
        if "level" in fieldnames:
            grade_key = "level"
        elif "DR_grade" in fieldnames:
            grade_key = "DR_grade"
        else:
            raise ValueError("CSV file must contain either 'level' or 'DR_grade' column.")
        
        has_quality=False
        if 'quality' in fieldnames:
            has_quality=True
        
        for row in reader:
            image_id = row['image'].strip()
            dr_val = row[grade_key].strip()
            if image_id[-4:]!= "jpeg" :
                # print(f"Warning: image_id {image_id} ,end:{image_id[-4:]} does not end with .jpeg, corrected")

                image_id+= '.jpeg'
            if not image_id or not dr_val:
                
                raise ValueError(f"Missing image or DR grade in row: {row}, row number: {reader.line_num}") 
            result[image_id] = {
                "classification":{
                    "text":dr_mapping[dr_val]
                }
            }
            if has_quality:
                q_val=row['quality'].strip()
                result[image_id]["classification"]['quality']=quality_mapping[q_val]
              
    return result


def merge_data_dict(train_official, test_custom):
    """
    合并两个由 CSV 解析后的数据字典：
      - 如果 test_custom 中的某个 image_name 在 train_official 中不存在，则将该键值对加入 train_official。
    返回合并后的字典。
    """
    merged = train_official.copy()
    for k, v in test_custom.items():
        if k not in merged:
            merged[k] = v
        else:
            for sub_k, sub_v in v.items():
                if sub_k not in merged[k]:
                    merged[k][sub_k] = sub_v
                else:
                    # 是否冲突
                    if merged[k][sub_k] != sub_v:
                        raise ValueError(f"Warning: Conflict for key {k}, sub-key {sub_k}. Keeping original value.")
    return merged

def gather_data(data_path, tar_path, prefix='eyepacs_', resize=(512,512)):
    """
    对新数据集进行预处理：
      - 数据集根目录下要求包含：
            trainLabels/trainLabels.csv
            Label_EyeQ_test.csv
            以及存放图像的 train 和 test 文件夹。
      - 分别解析 trainLabels 和测试 CSV，生成的标签字典 key 为图像文件名（含扩展名），值为诊断描述，
        然后将二者合并。
      - 遍历 train 和 test 文件夹，对每张图像调用 crop_resize_save 进行裁剪与 resize，
        并生成新的图像名称，格式为 "{prefix}_{原始文件名}.png"（保留原始文件名，仅修改后缀并加前缀）。
      - 同时生成一个 split 字典，记录各 split 下的处理后图像名称。
      - 最终生成的标注格式为：
            {
                new_image_name: {
                    "image_name": new_image_name,
                    "image_path": "images/{new_image_name}",
                    "diagnosis": {
                        "text": <diagnosis text>
                    }
                },
                ...
            }
      - 将标注信息保存到 tar_path/annotations.json，
        将 split 信息保存到 tar_path/split/official.json。
    
    参数：
        data_path (str): 新数据集根目录，要求包含 IMAGES 文件夹和相应 CSV 文件。
        tar_path (str): 预处理后数据存放目录。
        prefix (str): 处理后图片名称的前缀，默认为 "eyepacs_"
        resize (tuple): 目标尺寸，默认为 (512,512)
    
    返回：
        dict: 标注信息字典
    """
    # 创建目标目录
    os.makedirs(tar_path, exist_ok=True)
    images_out_dir = os.path.join(tar_path, "images")
    os.makedirs(images_out_dir, exist_ok=True)
    split_path = os.path.join(tar_path,'split.json')
    # 解析 CSV 文件生成标签字典
    train_csv_path = os.path.join(data_path, "trainLabels", "trainLabels.csv")
    test_csv_path = os.path.join(data_path, "Label_EyeQ_test.csv")
    train_eyeq_csv_path = os.path.join(data_path, "Label_EyeQ_train.csv")
    train_official = parse_csv(train_csv_path)
    test_custom = parse_csv(test_csv_path)
    train_eyeq=parse_csv(train_eyeq_csv_path)
    label_dict = merge_data_dict(train_official, test_custom)
    label_dict = merge_data_dict(label_dict, train_eyeq)

    final_annotations = {}
    official_split = {"train": [], "test": []}
    split_list = ["train", "test"]
    
    # 遍历每个 split 文件夹
    for split in split_list:
        image_src_dir = os.path.join(data_path, split)
        image_list=sorted(os.listdir(image_src_dir))
        total_image_num=len(image_list)
        for image_file in tqdm(image_list,desc=f'{split} process',unit='images',total=total_image_num):
            if not image_file.lower().endswith(('.png','.jpg','.jpeg','.bmp','.ppm')):
                continue
            # 保留原始文件名（不含扩展名），仅加前缀并统一保存为 PNG 格式
            base_name = os.path.splitext(image_file)[0]
            new_image_name = f"{prefix}_{base_name}.png"
            dest_image_path = os.path.join(images_out_dir, new_image_name)
            src_image_path = os.path.join(image_src_dir, image_file)
            # 进行图像预处理：裁剪、resize 并保存为 PNG
            crop_info=crop_resize_save(
                image_path = src_image_path,
                save_path = dest_image_path,
                resize = resize,
                crop_threshold = 15
            )
            # 获取对应的诊断描述（以原始文件名作为 key）
            diag_record = label_dict.get(image_file, "")
            final_annotations[new_image_name] = {
                "image_name": new_image_name,
                "image_path": dest_image_path,
                'original_path':src_image_path,
                'crop_info':crop_info,
                
            }
            if diag_record:
                final_annotations[new_image_name]["diagnosis"]=diag_record
            official_split[split].append(new_image_name)
    
    # 保存标注信息
    annotations_path = os.path.join(tar_path, "annotations.json")
    with open(annotations_path, "w", encoding="utf-8") as f:
        json.dump(final_annotations, f, indent=4)
    save_classification_stats(final_annotations,'./experiments/stat/eyepacs.json')
    
    # 保存分割信息
    with open(split_path, "w", encoding="utf-8") as f:
        json.dump(official_split, f, indent=4)
    
    return final_annotations

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Eyepacs 数据集预处理")
    parser.add_argument("--data_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_dataset/EyePACS",
                        help="新数据集根目录，要求其中包含 trainLabels, Label_EyeQ_test CSV 文件以及 train 和 test 图像文件夹")
    parser.add_argument("--tar_path", type=str, default="/home/zhangpinglu/data0/gy/Dataset/public_processed/EyePACS",
                        help="预处理后数据存放目录")
    parser.add_argument("--prefix", type=str, default="eyepacs",
                        help="处理后图片名称的前缀，默认为 'eyepacs'")
    parser.add_argument("--resize", type=int, nargs=2, default=[512,512],
                        help="目标尺寸，默认为 512 512")
    args = parser.parse_args()
    
    annotations = gather_data(args.data_path, args.tar_path, prefix=args.prefix, resize=tuple(args.resize))
    print("Preprocessing completed. Annotations saved.")
